import torch
import torch.nn as nn
import torchvision.models as models

# 编译通过，能正常输出，待结合实际调试；2025年3月25日

class SegNet(nn.Module):
    def __init__(self, num_classes):
        super(SegNet, self).__init__()

        # 编码器部分，使用 VGG16 的特征提取层
        vgg16 = models.vgg16(pretrained=True)
        self.encoder = nn.Sequential(*list(vgg16.features.children()))

        # 解码器部分
        self.decoder = nn.Sequential(
            # 第一层解码器
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(512, 512, kernel_size=2, stride=2),

            # 第二层解码器
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(256, 256, kernel_size=2, stride=2),

            # 第三层解码器
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(128, 128, kernel_size=2, stride=2),

            # 第四层解码器
            nn.Conv2d(128, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(64, 64, kernel_size=2, stride=2),

            # 第五层解码器
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, num_classes, kernel_size=3, padding=1)
        )

    def forward(self, x):
        # 编码器前向传播
        x = self.encoder(x)
        # 解码器前向传播
        x = self.decoder(x)
        return x

# 示例使用
if __name__ == "__main__":
    # 定义输入图像的大小和类别数量
    input_tensor = torch.randn(1, 3, 224, 224)
    num_classes = 10

    # 创建 SegNet 模型实例
    model = SegNet(num_classes)

    # 前向传播
    output = model(input_tensor)
    print("Output shape:", output.shape)